#include <scheme.h>
#include <linux/version.h>
#include <linux/module.h>

#include <linux/security.h>
#include <linux/security.h>
#include <linux/uaccess.h>  
#include <linux/kallsyms.h>
#include <linux/binfmts.h>
#include <linux/fs.h>
#include <linux/string.h>
#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0)
	#include <linux/lsm_hooks.h>
#endif
#include "init.h"
#include "data_type.h"
#include "policy_buffer.h"
#include "file_access.h"
#include "file_mon.h"
#include "globals.h"
 
#include "lsm.h"

extern int deep_debug;

#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0)
	struct security_hook_heads test_security_hook_heads;
#else
	static struct security_operations original_security_ops = {0}; 
	struct security_operations * g_ops = NULL;
	#define swap_security_ops(op)						\
		original_security_ops.op = g_ops->op; smp_wmb(); g_ops->op = os_##op;

	#define callback_security_ops(op)						\
		if(original_security_ops.op) g_ops->op = original_security_ops.op; smp_wmb();

#endif

static const char proc_name[] = "envset_proc";

static int test_mmap_file(struct file *file, unsigned long reqprot,
			     unsigned long prot, unsigned long flags)
{
	// -- printk("test-mmap-file......\n"); 
	return 0;
}

int os_bprm_set_creds(struct linux_binprm *bprm)
{
	int err = 0;
	os_security_t *sec = NULL;

	sec = (os_security_t *)kmalloc(sizeof(os_security_t), GFP_KERNEL);         // not free, memory leak !!!
	if (sec == NULL) {
		err = -ENOMEM;
		goto out;
	}
	memset(sec, 0, sizeof(os_security_t));

	if (strstr(bprm->filename, proc_name) != NULL) {
		sec->indicate = 1;
	}

	if (bprm->cred->security == NULL)
		bprm->cred->security = sec;
	else
		kfree(sec);

out:
	return err;
}

int os_mmap_file(struct file *file, unsigned long reqprot,
                unsigned long prot, unsigned long flags)
{
	os_security_t *sec = (os_security_t *)current->cred->security;
	if (file != NULL && sec != NULL && sec->indicate == 1) {
		printk("Process {%s} mmap file {%s}\n", proc_name, file->f_path.dentry->d_name.name);
		// printk("Process {%s} mmap file {%llx}\n", proc_name, file);
	}

	// read tick count
	if (sec != NULL && sec->indicate == 1) {
		uint64_t lo, hi;
		// RDTSC copies contents of 64-bit TSC into EDX:EAX
		__asm (
			"xor %%rax, %%rax\n\t"
			"cpuid\n\t"        
			"rdtsc\n\t" 
			: "=a" (lo), "=d" (hi)
		);
		uint64_t res = (hi << 32) | lo;
		printk("procname: %s \n", current->comm);
		printk("end time: %ld\n", res);
	}

	return 0;
}
 
int os_file_alloc_security(struct file *file)
{
	int rc = 0;
	#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0)
	
	#else	
		while (!original_security_ops.file_alloc_security);
		
		rc = original_security_ops.file_alloc_security(file);
		if(0 != rc){
			return rc;
		}
	#endif
	file_mon_alloc_security(file);
	return rc;
}

void os_file_free_security(struct file *file)
{
	file_mon_free_security(file);

	#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0)
	
	#else	
		while (!original_security_ops.file_free_security);
		original_security_ops.file_free_security(file);
	#endif
	return;
}
		
static int os_bprm_check_security (struct linux_binprm *bprm)
{
	int retval = 0;
 
	#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0)
	
	#else
		retval = original_security_ops.bprm_check_security(bprm);  
		if(0 != retval){
			return retval;
		}
	#endif 
	
	if(atomic_read(&g_globals.bWhiteListEnable)){
		if (pm_check_fs_integrity(bprm->file) != 0){
			printk(KERN_EMERG"file:%s is not in the whitelist!!!\n",bprm->filename);
			retval = -1;
		}
	}
		
	return retval;
}

static int os_file_permission(struct file * file, int mask)
{
	int rc = 0;

	#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0)
	
	#else
		while (!original_security_ops.file_permission);

		rc = original_security_ops.file_permission(file, mask);
		if(0 != rc){
			return rc;
		}
		
	#endif
		
	if(IS_ERR(file) || ( !file)){
		goto DONE;
	}

	rc = file_mon_permission_check(file, mask);
DONE: 
	return rc;
}

#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0)
	static struct security_hook_list test_hooks[] = {
		{ .head = &test_security_hook_heads.bprm_check_security,	\
				.hook = { .bprm_check_security = os_bprm_check_security } },
		{ .head = &test_security_hook_heads.file_permission,	\
				.hook = { .bprm_check_security = os_file_permission } },		
				
	};
#else
	
#endif

 
#define MAX_RO_PAGES 1024
static struct page *ro_pages[MAX_RO_PAGES];
static unsigned int ro_pages_len = 0;

static bool __init __test_ro_page(void *addr)
{
	unsigned int i;
	int unused;
	struct page *page;

	page = (struct page *) lookup_address((unsigned long) addr, &unused);
	if (!page)
		return -1;
	if (test_bit(_PAGE_BIT_RW, &(page->flags)))
		return 0;
	for (i = 0; i < ro_pages_len; i++)
		if (page == ro_pages[i])
			return 0;
	if (ro_pages_len == MAX_RO_PAGES)
		return -1;
	ro_pages[ro_pages_len++] = page;
	return 0;
}

#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0)
	static bool __check_pages(struct security_hook_heads *hooks)
	{
		int i;
		struct list_head *list = (struct list_head *) hooks;

		if (!probe_kernel_write(&list->next, list->next, sizeof(void *)))
			return 0;
		for (i = 0; i < ARRAY_SIZE(test_hooks); i++) {
			const unsigned int idx =
				((unsigned long) test_hooks[i].head
				 - (unsigned long) hooks)
				/ sizeof(struct list_head);
			struct list_head * __entry = &list[idx];
			struct list_head * __prev = __entry->prev;

			if (0 != __test_ro_page(&__prev->next) ||
				0 != __test_ro_page(&__entry->prev))
				return -1;
			if (!list_empty(__entry) &&
				0 != __test_ro_page(&list_last_entry
						  (__entry, struct security_hook_list,
						   list)->hook))
				return -1;
		}
		return 0;
	}
#endif

#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0)
	static inline void __hook_entry(struct security_hook_list * hook)
	{
		list_add_tail_rcu(&hook->list, hook->head);
	}
#endif

int lsm_hook(void)
{
	int retval = -EINVAL;
	 	
	#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0)
		struct security_hook_heads * hooks = (struct security_hook_heads *)kallsyms_lookup_name("security_hook_heads"); //-- ;// -- probe_security_hook_heads();
	#else
		struct security_operations ** old = (struct security_operations **)kallsyms_lookup_name("security_ops");
		
	#endif
	
	#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0)
		if (!hooks){
			printk("no security hook heads \n");
			goto DONE;
		}

		for (idx = 0; idx < ARRAY_SIZE(test_hooks); idx++){
			test_hooks[idx].head = ((void *) hooks)
				+ ((unsigned long) test_hooks[idx].head)
				- ((unsigned long) &test_security_hook_heads);
		}
		
		if (0 != __check_pages(hooks)) {
			return retval;
		}
	 
		for (idx = 0; idx < ro_pages_len; idx++){
			set_bit(_PAGE_BIT_RW, &(ro_pages[idx]->flags));
		}
		
		for (idx = 0; idx < ARRAY_SIZE(test_hooks); idx++){
			__hook_entry(&test_hooks[idx]);
		}
			 
		for (idx = 0; idx < ro_pages_len; idx++){
			clear_bit(_PAGE_BIT_RW, &(ro_pages[idx]->flags));
		}		
		
	#else
		if(!old){
			goto DONE;
		}
		
		g_ops = *old;
		
		swap_security_ops(bprm_check_security);
		swap_security_ops(file_permission);
		swap_security_ops(file_alloc_security);
		swap_security_ops(file_free_security);
		swap_security_ops(bprm_set_creds);
		swap_security_ops(mmap_file);
	#endif 
	
	retval = 0;	
DONE:	
	return retval;
}

void lsm_unhook(void)
{
	#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0)
		int idx;
	#endif
	
	#if LINUX_VERSION_CODE >= KERNEL_VERSION(4, 2, 0)
		for (idx = 0; idx < ro_pages_len; idx++){
			set_bit(_PAGE_BIT_RW, &(ro_pages[idx]->flags));
		}
		
		for (idx = 0; idx < ARRAY_SIZE(test_hooks); idx++){
			list_del_rcu(&test_hooks[idx].list);
		}	
		
		for (idx = 0; idx < ro_pages_len; idx++){
			clear_bit(_PAGE_BIT_RW, &(ro_pages[idx]->flags));
		} 
	#else 
		callback_security_ops(file_permission);
		callback_security_ops(bprm_check_security);
		callback_security_ops(file_alloc_security);
		callback_security_ops(file_free_security);
		callback_security_ops(mmap_file);
		callback_security_ops(bprm_set_creds);
	#endif	
	return ;
}	
	
